/*
 * gmx_elasticnet.c
 *
 *  Created on: Jan 26, 2011
 *      Author: marchi
 */
/*
 *
 *                This source code is part of
 *
 *                 G   R   O   M   A   C   S
 *
 *          GROningen MAchine for Chemical Simulations
 *
 * Written by David van der Spoel, Erik Lindahl, Berk Hess, and others.
 * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
 * Copyright (c) 2001-2009, The GROMACS development team,
 * check out http://www.gromacs.org for more information.

 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * If you want to redistribute modifications, please consider that
 * scientific software is very special. Version control is crucial -
 * bugs must be traceable. We will be happy to consider code for
 * inclusion in the official distribution, but derived work must not
 * be called official GROMACS. Details are found in the README & COPYING
 * files - if they are missing, get the official version at www.gromacs.org.
 *
 * To help us fund GROMACS development, we humbly ask that you cite
 * the papers on the package - you can find them in the top README file.
 *
 * For more info, check our website at http://www.gromacs.org
 */
#include "copyrite.h"
#include "filenm.h"
#include "macros.h"
#include "pbc.h"
#include "smalloc.h"
#include "statutil.h"
#include "vec.h"
#include "xvgr.h"
#include "typedefs.h"

#include "nbsearch.h"
#include "trajana.h"
#include "futil.h"
#include "index.h"
#include "strdb.h"
#include "string2.h"
#include "pbc.h"
#include "rmpbc.h"
#include "elastic.h"
#include "mkl.h"
#include "gmx_fatal.h"
#include "do_fit.h"

static double **tr;

static matrix ** matrirj;
static double ** matdist;
static rvec * avgri;
static rvec * rri;
FILE * myfp;
static t_block ref;
static real * w_lst;
static rvec * xrefs;
static int NatomProtein=0;


static void put_molecule_com_in_box(int unitcell_enum,int ecenter,
                                    t_block *mols,
                                    int natoms,t_atom atom[],
                                    int ePBC,matrix box,rvec x[])
{
    atom_id i,j;
    int     d;
    rvec    com,new_com,shift,box_center;
    real    m;
    double  mtot;
    t_pbc   pbc;
    calc_box_center(ecenter,box,box_center);
    set_pbc(&pbc,ePBC,box);
    if (mols->nr <= 0)
        gmx_fatal(FARGS,"There are no molecule descriptions. I need a tpr file for this pbc option.");
    for(i=0; (i<mols->nr); i++) {
        /* calc COM */
        clear_rvec(com);
        mtot = 0;
        for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
            m = atom[j].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[j][d];
            mtot += m;
        }
        /* calculate final COM */
        svmul(1.0/mtot, com, com);
	for(d=0; d<DIM; d++)
        /* check if COM is outside box */
        copy_rvec(com,new_com);
        switch (unitcell_enum) {
        case euRect:
            put_atoms_in_box(box,1,&new_com);
            break;
        case euTric:
            put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
            break;
        case euCompact:
            put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
            break;
        }
        if (debug) fprintf(debug," My debug %d %d \n",ePBC,ecenter);
        rvec_sub(new_com,com,shift);
        if (norm2(shift) > 0) {
            if (debug)
                fprintf (debug,"\nShifting position of molecule %d "
                         "by %8.3f  %8.3f  %8.3f\n", i+1, PR_VEC(shift));
            for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
                rvec_inc(x[j],shift);
            }
        }else if(debug) fprintf (debug,"\n My Debug 2 %d  %f %f %f \n", i+1, new_com[XX], new_com[YY], new_com[ZZ]);
    }
}
static void put_residue_com_in_box(int unitcell_enum,int ecenter,
                                   int natoms,t_atom atom[],
                                   int ePBC,matrix box,rvec x[])
{
    atom_id i, j, res_start, res_end, res_nat;
    int     d, presnr;
    real    m;
    double  mtot;
    rvec    box_center,com,new_com,shift;
    calc_box_center(ecenter,box,box_center);
    presnr = NOTSET;
    res_start = 0;
    clear_rvec(com);
    mtot = 0;
    for(i=0; i<natoms+1; i++) {
#ifdef OLDVERSION
        if (i == natoms || (presnr != atom[i].resnr && presnr != NOTSET)) {
#else
        if (i == natoms || (presnr != atom[i].resind && presnr != NOTSET)) {
#endif
            /* calculate final COM */
            res_end = i;
            res_nat = res_end - res_start;
            svmul(1.0/mtot, com, com);
            /* check if COM is outside box */
            copy_rvec(com,new_com);
            switch (unitcell_enum) {
            case euRect:
                put_atoms_in_box(box,1,&new_com);
                break;
            case euTric:
                put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
                break;
            case euCompact:
                put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
                break;
            }
            rvec_sub(new_com,com,shift);
            if (norm2(shift)) {
                if (debug)
#ifdef OLDVERSION
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resnr+1,
#else
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resind+1,
#endif
                             res_start+1, res_end+1, PR_VEC(shift));
                for(j=res_start; j<res_end; j++)
                    rvec_inc(x[j],shift);
            }
            clear_rvec(com);
            mtot = 0;
            /* remember start of new residue */
            res_start = i;
        }
        if (i < natoms) {
            /* calc COM */
            m = atom[i].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[i][d];
            mtot += m;
#ifdef OLDVERSION
            presnr = atom[i].resnr;
#else
            presnr = atom[i].resind;
#endif
        }
    }
}
static void center_x(int ecenter,rvec x[],matrix box,int n,int nc,atom_id ci[]){
    int i,m,ai;
    rvec cmin,cmax,box_center,dx;
    if (nc > 0) {
        copy_rvec(x[ci[0]],cmin);
        copy_rvec(x[ci[0]],cmax);
        for(i=0; i<nc; i++) {
            ai=ci[i];
            for(m=0; m<DIM; m++) {
                if (x[ai][m] < cmin[m])
                    cmin[m]=x[ai][m];
                else if (x[ai][m] > cmax[m])
                    cmax[m]=x[ai][m];
            }
        }
        calc_box_center(ecenter,box,box_center);
        for(m=0; m<DIM; m++)
            dx[m] = box_center[m]-(cmin[m]+cmax[m])*0.5;
        for(i=0; i<n; i++)
            rvec_inc(x[i],dx);
    }
}
static void acc_x(int ncent, int *cindex, rvec *x_in, rvec *x_out){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x_out[i][XX]+=x_in[i][XX];
		x_out[i][YY]+=x_in[i][YY];
		x_out[i][ZZ]+=x_in[i][ZZ];
	}
}
static void avg_x(int ncent, int *cindex,int n, rvec *x, rvec  add){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x[i][XX]/=(real) n;
		x[i][YY]/=(real) n;
		x[i][ZZ]/=(real) n;
		x[i][XX]+=add[XX];
		x[i][YY]+=add[YY];
		x[i][ZZ]+=add[ZZ];
	}
}

struct gmx_residuetype
{
    int      n;
    char **  resname;
    char **  restype;

};


/*! \brief
 * Template analysis data structure.
 */
typedef struct
{
    gmx_ana_selection_t *refsel;
    FILE                *fp;
    real                *ave;
    real                *n;
    gmx_ana_nbsearch_t  *nb;
} t_analysisdata;



/*! \brief
 * Function that does the analysis for a single frame.
 *
 * It is called once for each frame.
 */
static gmx_bool gmx_residuetype_is_hydrophile(gmx_residuetype_t rt, const char *resnm)
{
    gmx_bool rc;
    const char *p_type;

    if(gmx_residuetype_get_type(rt,resnm,&p_type)==0 &&
       gmx_strcasecmp(p_type,"hydrophile")==0)
    {
        rc=TRUE;
    }
    else
    {
        rc=FALSE;
    }
    return rc;
}

static gmx_bool gmx_residuetype_is_hydrophobe(gmx_residuetype_t rt, const char *resnm)
{
    gmx_bool rc;
    const char *p_type;

    if(gmx_residuetype_get_type(rt,resnm,&p_type)==0 &&
       gmx_strcasecmp(p_type,"hydrophobe")==0)
    {
        rc=TRUE;

    }
    else
    {
        rc=FALSE;
    }
    return rc;
}

static void create_residues_atindex(resi_t * res, int nresidue, int natoms, t_atoms * atoms)
{

    atom_id i, res_start, res_end, res_nat;
    int     d,presnr,o;
    presnr = NOTSET;
    res_start = 0;
    d=0;
    printf("\n");
    for(i=0; i<natoms+1; i++) {
    	if (i == natoms || (presnr != atoms->atom[i].resind && presnr != NOTSET)) {
    		res_end = i;
            res_nat = res_end - res_start;
            res[d].nr=res_nat;
            snew(res[d].index,res_nat);
            double mtot=0.0;
            for(o=res_start;o<res_end;o++) {
            	res[d].index[o-res_start]=o;
            	mtot+=atoms->atom[o].m;
            }
            res[d].mtot=mtot;
            res[d].name=*atoms->resinfo[atoms->atom[i].resind].name;

            /* remember start of new residue */
            res_start = i;
            d++;
        }
        if (i < natoms) {
        	presnr = atoms->atom[i].resind;
        }

    }
}

static void create_residue_cm(gmx_ana_traj_t *d,residue_t * restot, residue_t * resprot){
    char    line[STRLEN];
    char    resname[STRLEN],restype[STRLEN],dum[STRLEN];

	FILE *db;
    struct gmx_residuetype *rt,*rtype;

    snew(rt,1);
    rt->n        = 0;
    rt->resname  = NULL;
    rt->restype = NULL;

	db=libopen("residuetypes.dat");

    while(get_a_line(db,line,STRLEN))
    {
    	strip_comment(line);
    	trim(line);
    	if(line[0]!='\0')
    	{
    		if(sscanf(line,"%s %s %s",resname,restype,dum)!=2)
    		{
    			gmx_fatal(FARGS,"Incorrect number of columns (2 expected) for line in hydrophobic.dat");
    		}
    		gmx_residuetype_add(rt,resname,restype);
    	}
    }

    fclose(db);

    snew(rtype,1);
    rtype->n        = 0;
    rtype->resname  = NULL;
    rtype->restype = NULL;

	db=libopen("hydrophobic.dat");

	while(get_a_line(db,line,STRLEN))
    {
    	strip_comment(line);
    	trim(line);
    	if(line[0]!='\0')
    	{
    		if(sscanf(line,"%s %s %s",resname,restype,dum)!=2)
    		{
    			gmx_fatal(FARGS,"Incorrect number of columns (2 expected) for line in hydrophobic.dat");
    		}
    		gmx_residuetype_add(rtype,resname,restype);
    	}
    }

    fclose(db);
	t_topology * top;
	gmx_ana_get_topology(d, FALSE, &top, NULL);

    t_atoms * atoms;
    static resi_t ** resid, * resid_a;

    int natoms=top->atoms.nr;

    int nresidue=top->atoms.nres;
    snew(resid_a,nresidue);
    atoms=&top->atoms;
    create_residues_atindex(resid_a, nresidue, natoms, atoms);

	const char * p_restype;
    int q=0;
	for(int i=0;i<nresidue;i++){
		gmx_residuetype_get_type(rt,*atoms->resinfo[i].name, &p_restype);
		printf(" %s \n",p_restype);
		if(gmx_strcasecmp(p_restype,"Protein")){
			break;
		}
		q++;
	}
    snew(resid,q);
    for(int i=0;i<q;i++) resid[i]=&resid_a[i];

    restot->nr=nresidue;
    resprot->nr=q;
    restot->list=resid_a;
    resprot->list=*resid;
    NatomProtein=resprot->list[resprot->nr-1].index[resprot->list[resprot->nr-1].nr-1];
    snew(w_lst,NatomProtein);
    snew(xrefs,NatomProtein);
}

static int analyze_frame(t_topology *top, t_trxframe *fr, t_pbc *pbc,
              int nr, gmx_ana_selection_t *sel[], void *data)
{
    /* Here, you can do whatever analysis your program requires for a frame. */

	static int nframes=0;
	residue_t * resprot=(residue_t *)data;

    t_atoms * atoms=&top->atoms;
	int ecenter=0;
    int ePBC;
	const char * warn;
	int unitcell_enum=euCompact;
    gmx_rmpbc_t  gpbc=NULL;

    //    int ru=gmx_ana_get_topconf(trj,NULL, NULL, &ePBC);
    ePBC=fr->ePBC;

	set_pbc(pbc,ePBC,fr->box);

	gpbc = gmx_rmpbc_init(&top->idef,ePBC,fr->natoms,fr->box);

    gmx_rmpbc_trxfr(gpbc,fr);

/*
	warn = put_atoms_in_compact_unitcell(ePBC,ecenter,fr->box,
			fr->natoms,fr->x);
*/
	put_molecule_com_in_box(unitcell_enum,ecenter,
				&top->mols,fr->natoms,top->atoms.atom,ePBC,fr->box,fr->x);

	reset_x(NatomProtein, NULL, NatomProtein, NULL,fr->x,w_lst);
    if(!nframes){
    	copy_rvecn(fr->x,xrefs,0,NatomProtein);
    }
    do_fit(NatomProtein,w_lst,xrefs,fr->x);


	rvec cm;
	rvec * cmi;
	snew(cmi,resprot->nr);
	for(int i=0;i<resprot->nr;i++){
		double mass=resprot->list[i].mtot;
		for(int o=0;o<DIM;o++)cm[o]=0.0;
		for(int p=0;p<resprot->list[i].nr;p++){
			int q=resprot->list[i].index[p];
			real m=atoms->atom[q].m/mass;
			for(int o=0;o<DIM;o++)cm[o]+=(fr->x[q][o])*m;
		}
		for(int o=0;o<DIM;o++) cmi[i][o]=cm[o];
	}
	rvec dx;
	fprintf(myfp," Frame No. = %d \n",nframes);
	for(int i=0;i<resprot->nr;i++){
		for(int o=0;o<DIM; o++){
			avgri[i][o]+=cmi[i][o];
			rri[i][o]+=cmi[i][o]*cmi[i][o];
		}
		for(int j=i; j<resprot->nr;j++){
			pbc_dx(pbc,cmi[i],cmi[j],dx);
			double rsp=norm(dx);
			matdist[i][j]+=rsp;
			for(int o=0;o<DIM; o++){
				for(int p=0;p<DIM;p++){
					matrirj[i][j][o][p]+=cmi[i][o]*cmi[j][p];
				}
			}
		}
	}
	if (gpbc != NULL)
    {
        gmx_rmpbc_done(gpbc);
    }


    /* We need to return 0 to tell that everything went OK */
	nframes++;
    return 0;
}

static void postprocessing(gmx_ana_traj_t *d,residue_t * resprot){

	t_topology * top;
	gmx_ana_get_topology(d, FALSE, &top, NULL);

	real ** mat;
	snew(mat,resprot->nr*DIM);
	for(int i=0;i<resprot->nr*DIM;i++) snew(mat[i],resprot->nr*DIM);

	int nframes;

	if(gmx_ana_get_nframes(d, &nframes)){
        gmx_fatal(FARGS,"This is impossible this function cannot return 1!!");
	}
	rvec * cmi;
	snew(cmi,resprot->nr);
	rvec * cmri;
	snew(cmri,resprot->nr);
	real ** matdist0;
	snew(matdist0,resprot->nr);
	for(int i=0;i<resprot->nr;i++) snew(matdist0[i],resprot->nr);

	for(int i=0;i<resprot->nr;i++){
		real rms=0.0;
		for(int o=0;o<DIM; o++){
			cmi[i][o]=avgri[i][o]/(real) nframes;
			cmri[i][o]=rri[i][o]/(real) nframes;
		}
		for(int o=0;o<DIM; o++){
			rms+=cmri[i][o]-cmi[i][o]*cmi[i][o];
		}
		printf("%d %f \n",i,rms);
	}

	int ia,ja;
	for(int i=0;i<resprot->nr;i++){
		for(int j=i; j<resprot->nr;j++){
			matdist0[i][j]=matdist[i][j]/(real) nframes;
			for(int o=0;o<DIM; o++){
				ia=DIM*i+o;
				for(int p=0;p<DIM;p++){
					ja=DIM*j+p;
					mat[ia][ja]=matrirj[i][j][o][p]/(real) nframes-cmi[i][o]*cmi[j][p];
				}
			}
			if(i==j) {
				double rsp=mat[DIM*i][DIM*j]+mat[DIM*i+1][DIM*j+1]+mat[DIM*i+2][DIM*j+2];
				printf(" rms 2 %d %f \n",j,rsp);
			}
		}

	}

}
/*! \brief
 * Function that implements the analysis tool.
 *
 * Following the style of Gromacs analysis tools, this function is called
 * \p gmx_something.
 */
int
gmx_elasticnet(int argc, char *argv[])
{
    const char         *desc[] = {
        "This is a template for writing your own analysis tools for",
        "Gromacs. The advantage of using Gromacs for this is that you",
        "have access to all information in the topology, and your",
        "program will be able to handle all types of coordinates and",
        "trajectory files supported by Gromacs. In addition,",
        "you get a lot of functionality for free from the trajectory",
        "analysis library, including support for flexible dynamic",
        "selections. Go ahead an try it![PAR]",
        "To get started with implementing your own analysis program,",
        "follow the instructions in the README file provided.",
        "This template implements a simple analysis programs that calculates",
        "average distances from the a reference group to one or more",
        "analysis groups.",
    };

    /* Command-line arguments */
    real                cutoff = 0;
    gmx_bool                bArg   = FALSE;
    t_pargs             pa[] = {
        {"-cutoff", FALSE, etREAL, {&cutoff},
         "Cutoff for distance calculation (0 = no cutoff)"},
        {"-arg2",   FALSE, etBOOL, {&bArg},
         "Example argument 2"},
    };
    /* The second argument is for demonstration purposes only */

    /* Output files */

    t_filenm            fnm[] = {
        {efDAT, "-o", "Forces", ffOPTWR},
    };
#define NFILE asize(fnm)

    gmx_ana_traj_t       *trj;
    gmx_ana_selection_t **sel;
    output_env_t          oenv;


    CopyRight(stderr, argv[0]);
    /* Here, we can use flags to specify requirements for the selections and/or
     * other features of the library. */
    gmx_ana_traj_create(&trj, ANA_REQUIRE_TOP);

    /* If required, other functions can also be used to configure the library
     * before calling parse_trjana_args(). */

    parse_trjana_args(trj, &argc, argv, PCA_CAN_VIEW,
                      NFILE, fnm, asize(pa), pa, asize(desc), desc, 0, NULL,
                      &oenv);

    /* You can now do any initialization you wish, using the information
     * from the trj structure.
     * In particular, you should store any command-line parameter values that
     * the analysis part requires into d for them to be accessible. */

    /* We then allocate memory in d to store intermediate data
     * and initialize the counters to zero */

    /* Now, we do the actual analysis */
    residue_t restot, resprot;

	myfp=fopen("pippo.dat","w");
    create_residue_cm(trj,&restot,&resprot);

	int  ngprs=0;

	if(!gmx_ana_get_ngrps(trj, &ngprs)) gmx_ana_get_anagrps(trj, &sel);

	if(ngprs != 1)
		gmx_fatal(FARGS," Only one group must be selected!! ");

	ref.nr = sel[0]->p.m.nr;
	if(ref.nr == 0)
		gmx_fatal(FARGS," No groups selected cannot run. ");

    snew(ref.index,ref.nr);
	for (int i = 0; i < ref.nr; ++i) ref.index[i]=sel[0]->p.m.orgid[i];
    for (int i = 0; i < ref.nr; ++i) printf(" %d", ref.index[i]);
    for (int i = 0;i<ref.nr; ++i) w_lst[ref.index[i]]=1.0/(real) ref.nr;


	snew(matrirj,resprot.nr);
	snew(matdist,resprot.nr);
	snew(avgri,resprot.nr);
	snew(rri,resprot.nr);
	for(int i=0;i<resprot.nr;i++){
		snew(matrirj[i],resprot.nr);
		snew(matdist[i],resprot.nr);
	}

    gmx_ana_do(trj, 0, &analyze_frame, &resprot);

    postprocessing(trj,&resprot);

    /* Now, the analysis has been done for all frames, and you can access the
     * results in d. Here, you should post-process your data and write out any
     * averaged properties. */

    /* Here, we could free some memory, but this usually not necessary as we
     * are going to quit anyways. */

    thanx(stderr);
    return 0;
}
